/*
 * Copyright 2014 Щекотова Ангелина Александровна.
 * Этот файл — часть J4JClientServer Framework
 * J4JClientServer Framework - свободная программа: вы можете перераспространять ее и/или изменять
 * ее на условиях Меньшей Стандартной общественной лицензии GNU в том виде, в каком
 * она была опубликована Фондом свободного программного обеспечения; либо
 * версии 3 лицензии, либо (по вашему выбору) любой более поздней версии.
 *
 * J4JClientServer Framework распространяется в надежде, что она будет полезной,
 * но БЕЗО ВСЯКИХ ГАРАНТИЙ; даже без неявной гарантии ТОВАРНОГО ВИДА
 * или ПРИГОДНОСТИ ДЛЯ ОПРЕДЕЛЕННЫХ ЦЕЛЕЙ. Подробнее см. в Меньшей Стандартной
 * общественной лицензии GNU.
 *
 * Вы должны были получить копию Меньшей Стандартной общественной лицензии GNU
 * вместе с этой программой. Если это не так, см.
 * <http://www.gnu.org/licenses/>.
 */

package ru.java4java.security;

import org.apache.log4j.Logger;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.*;
import java.math.BigInteger;
import java.security.*;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPrivateKeySpec;
import java.security.spec.RSAPublicKeySpec;

public class CipherEngine {

    // Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());

    private transient static final  Logger log = Logger.getLogger(CipherEngine.class);

    private final String algorithm;
    private final String provider;

    private final KeyPairGenerator keyPairGenerator;
    private final KeyFactory keyFactory;
    private final Cipher cipher;
    private KeyPair keyPair;

    private static boolean init = false;

    public CipherEngine(String algorithm, String provider , int keySize) {
        try {
            this.algorithm = algorithm;
            this.provider = provider;
            keyPairGenerator = KeyPairGenerator.getInstance(algorithm, provider);
            keyFactory = KeyFactory.getInstance(algorithm, provider);
            cipher = Cipher.getInstance(algorithm, provider);
        } catch (NoSuchAlgorithmException e) {
            log.fatal("NoSuchAlgorithmException", e);
            throw new RuntimeException(e);
        } catch (NoSuchProviderException e) {
            log.fatal("NoSuchProviderException", e);
            throw new RuntimeException(e);
        } catch (NoSuchPaddingException e) {
            log.fatal("NoSuchPaddingException", e);
            throw new RuntimeException(e);
        }
        generateNewKeys(keySize);
    }

    public String getAlgorithm() {
        return algorithm;
    }

    public String getProvider() {
        return provider;
    }

    public void generateNewKeys(int keySize) {
        keyPairGenerator.initialize(keySize); //1024 used for normal securities
        keyPair = keyPairGenerator.generateKeyPair();
    }

    public void setKeys(PublicKey publicKey, PrivateKey privateKey) { keyPair = new KeyPair(publicKey, privateKey); }

    public PublicKey getPublicKey() { return keyPair.getPublic(); }

    public PrivateKey getPrivateKey() { return keyPair.getPrivate(); }

    public byte[] encryptData(final byte[] data, PublicKey pubKey) throws EncryptDataException {
        byte[] encryptedData;

        try {
            cipher.init(Cipher.ENCRYPT_MODE, pubKey);
        } catch (InvalidKeyException e) {
            log.error("InvalidKeyException", e);
            throw new EncryptDataException("InvalidKeyException", e);
        }

        try {
            encryptedData = cipher.doFinal(data);
        } catch (IllegalBlockSizeException e) {
            log.error("IllegalBlockSizeException", e);
            throw new EncryptDataException("IllegalBlockSizeException", e);
        } catch (BadPaddingException e) {
            log.error("BadPaddingException", e);
            throw new EncryptDataException("BadPaddingException", e);
        }

        return encryptedData;
    }

    public byte[] decryptData(byte[] data, PrivateKey privateKey) throws DecryptDataException {
        byte[] descryptedData = null;

        try {
            cipher.init(Cipher.DECRYPT_MODE, privateKey);
        } catch (InvalidKeyException e) {
            log.error("InvalidKeyException", e);
            throw new DecryptDataException("InvalidKeyException", e);
        }

        try {
            descryptedData = cipher.doFinal(data);
        } catch (IllegalBlockSizeException e) {
            log.error("IllegalBlockSizeException", e);
            throw new DecryptDataException("IllegalBlockSizeException", e);
        } catch (BadPaddingException e) {
            log.error("BadPaddingException", e);
            throw new DecryptDataException("BadPaddingException", e);
        }
        return descryptedData;
    }

    public void writeKeys(OutputStream outputStream, BigInteger mod, BigInteger exp) throws IOException {
        ObjectOutputStream oos;
        oos = new ObjectOutputStream(outputStream);
        oos.writeObject(mod);
        oos.writeObject(exp);
    }

    public PublicKey readPublicKey(InputStream inputStream) throws IOException, InvalidKeySpecException {
        Object[] objects = readObjects(inputStream, 2);
        RSAPublicKeySpec rsaPublicKeySpec = new RSAPublicKeySpec(
          (BigInteger)objects[0], /* modulus  */
          (BigInteger)objects[1]  /* exponent */
        );
        return keyFactory.generatePublic(rsaPublicKeySpec);
    }

    public PrivateKey readPrivateKey(InputStream inputStream) throws IOException, InvalidKeySpecException {
        Object[] objects = readObjects(inputStream, 2);
        RSAPrivateKeySpec rsaPrivateKeySpec = new RSAPrivateKeySpec(
          (BigInteger)objects[0], /* modulus  */
          (BigInteger)objects[1]  /* exponent */
        );
        return keyFactory.generatePrivate(rsaPrivateKeySpec);
    }

    private Object[] readObjects(InputStream inputStream, int objectCount)
      throws IOException
    {
        Object[] objects = new Object[objectCount];
        ObjectInputStream ois = new ObjectInputStream(inputStream);
        try {
            for (int i = 0; i < objectCount; i++) {
                objects[i] = ois.readObject();
            }
        } catch (ClassNotFoundException e) {
            log.fatal("ClassNotFoundException", e);
            throw new RuntimeException(e);
        }
        return objects;
    }
}
